from faker import Faker
import time
import requests
import m3u8
# 报错异常，则手动至包的位置，将目录 crypto 改为 Crypto
from Crypto.Util.Padding import pad
from Crypto.Cipher import AES

# 模拟数据
fk = Faker(locale="zh_CN")
# 设置请求头信息
headers = {
    'User-Agent': fk.user_agent()
}


# 获取实际地址
def get_real_url(url):
    playlist = m3u8.load(uri=url, headers=headers)
    return playlist.playlists[0].absolute_uri


# 获取解密 key
def get_encrypt_key(playlist):
    key_url = playlist.keys[-1].absolute_uri
    key = requests.get(key_url, headers=headers, timeout=3000).content
    print("获取解密 Key ...")
    return key


# 解密
def decrypt(data, key):
    data = pad(data_to_pad=data, block_size=AES.block_size)
    aes = AES.new(key=key, mode=AES.MODE_CBC, iv=key)
    return aes.decrypt(data)


def download(url, save, encrypt=False):
    # 播放地址
    playlist = m3u8.load(uri=get_real_url(url), headers=headers)
    # 解密
    key = get_encrypt_key(playlist) if encrypt else ''

    size = 0
    start = time.time()
    total = len(playlist.segments)
    print("开始下载...")
    for i, seg in enumerate(playlist.segments, 1):
        r = requests.get(seg.absolute_uri, headers=headers, timeout=3000)
        data = r.content
        # 解密
        data = decrypt(data, key) if encrypt else data
        # 数据大小
        size += len(data)
        with open(save, "ab" if i != 1 else "wb") as f:
            f.write(data)
        print(
            f"\r下载进度({i}/{total})，已下载：{size/1024/1024:.2f}MB，下载已耗时：{time.time()-start:.2f}s", end=" ")


url = 'xxx/index.m3u8'
download(url, './result.mp4', True)
